"use client"

import * as ToggleGroupPrimitive from "@radix-ui/react-toggle-group"
import type { VariantProps } from "class-variance-authority"
import * as React from "react"

import { ny } from "~/lib/utils"
import { toggleVariants } from "~/registry/default/ui/toggle"

const ToggleGroupContext = React.createContext<
   VariantProps<typeof toggleVariants>
>({
   size: "default",
   variant: "default",
})

const ToggleGroup = React.forwardRef<
   React.ElementRef<typeof ToggleGroupPrimitive.Root>,
   React.ComponentPropsWithoutRef<typeof ToggleGroupPrimitive.Root> &
      VariantProps<typeof toggleVariants>
>(({ className, variant, size, children, ...props }, ref) => (
   <ToggleGroupPrimitive.Root
      ref={ref}
      className={ny("flex items-center justify-center gap-1", className)}
      {...props}
   >
      <ToggleGroupContext.Provider value={{ variant, size }}>
         {children}
      </ToggleGroupContext.Provider>
   </ToggleGroupPrimitive.Root>
))

ToggleGroup.displayName = ToggleGroupPrimitive.Root.displayName

const ToggleGroupItem = React.forwardRef<
   React.ElementRef<typeof ToggleGroupPrimitive.Item>,
   React.ComponentPropsWithoutRef<typeof ToggleGroupPrimitive.Item> &
      VariantProps<typeof toggleVariants>
>(({ className, children, variant, size, ...props }, ref) => {
   const context = React.useContext(ToggleGroupContext)

   return (
      <ToggleGroupPrimitive.Item
         ref={ref}
         className={ny(
            toggleVariants({
               variant: context.variant || variant,
               size: context.size || size,
            }),
            className,
         )}
         {...props}
      >
         {children}
      </ToggleGroupPrimitive.Item>
   )
})

ToggleGroupItem.displayName = ToggleGroupPrimitive.Item.displayName

export { ToggleGroup, ToggleGroupItem }
